import traceback
import base64
import os
import shutil
import glob
import zipfile
import magic
import time
import pymisp
from threading import Thread, Lock
from typing import Union, Sequence, List, Tuple, Optional
from pathlib import Path
from urllib3.exceptions import ProtocolError
from androidsandbox.device import Device, Devices, DeviceError, DockerDevice
from androidsandbox.utils import logger, Utils, CFG,fail_apk_logger


class DeviceThread(Thread):
    """Each device is controlled by one thread."""
    def __init__(self, sample_count: List, device: Device,
                 misp_tags: List[str], lock: Lock):
        super().__init__()
        self.sample_count = sample_count
        self.device = device
        self.misp = Utils.init_misp(CFG['misp']['url'], CFG['misp']['key'])
        self.misp_tags = {'or': [], 'and': [], 'not': []}
        self.download_lock = lock
        self.host_dir = None
        self.update_host_dir()
        # for building build_complex_query
        if misp_tags:
            if isinstance(misp_tags, list):
                misp_tags = misp_tags
            else:
                misp_tags = [misp_tags]
        else:
            misp_tags = []
        for tag in misp_tags:
            if tag.startswith('or:'):
                self.misp_tags['or'].append(tag[3:])
            elif tag.startswith('and:'):
                self.misp_tags['and'].append(tag[4:])
            elif tag.startswith('not:'):
                self.misp_tags['not'].append(tag[4:])

    def update_host_dir(self):
        if isinstance(self.device, DockerDevice):
            self.host_dir = self.device.super_host_dir
        else:
            self.host_dir = self.device.host_dir

    def download_one_sample(self) -> Tuple[Optional[str], Optional[int],
                                           Optional[str], Optional[str]]:
        """
         Search attributes of samples to be tested.
        :return: (event_id, sample_md5, event_uuid, package name)
        """
        try:
            with self.download_lock:
                res = self.misp.search(controller='events',
                                       tags=self.misp.build_complex_query(
                                           and_parameters=self.misp_tags['and'],
                                           not_parameters=self.misp_tags['not'],
                                           or_parameters=self.misp_tags['or']),
                                       type_attribute='threat-actor',
                                       **{'value': '%package-name%'},
                                       with_attachments=True,
                                       limit=1)

                if not res:
                    return None, None, None, None

                event = res[0]['Event']
                event_id = int(event['id'])
                event_uuid = event['uuid']
                sample_md5, zip_data, package_name = None, None, None

                for obj in event['Object']:
                    for attr in obj['Attribute']:
                        if attr['type'] == 'malware-sample':
                            zip_data = attr['data']
                        elif attr['type'] == 'md5':
                            sample_md5 = attr['value']

                for attr in event['Attribute']:
                    if attr['value'].startswith('package-name'):
                        package_name = attr['value'].split('"')[1]
                        break

                if not (sample_md5 and zip_data and package_name):
                    raise ValueError("Required sample attributes are missing.")
                self.misp.tag(event_uuid, 'secheart:analysis="ui-testing"')
                self.misp.tag(event_uuid, 'secheart:analysis="fail-apk"')
                time.sleep(3)

            zip_path = os.path.join(self.host_dir, f'{sample_md5}.zip')

            try:
                with open(zip_path, 'wb') as f:
                    f.write(base64.b64decode(zip_data))
            except IOError as e:
                raise DeviceError(f'Error writing file {zip_path}: {e}')

            try:
                with zipfile.ZipFile(zip_path, 'r') as zf:
                    zf.extractall(self.host_dir, pwd=b'infected')
            except zipfile.BadZipFile:
                raise DeviceError(f'Error unzipping file {zip_path}')

            os.remove(zip_path)

            sample_file_name = None
            for p in glob.glob(os.path.join(self.host_dir, '*')):
                if zipfile.is_zipfile(p):
                    with zipfile.ZipFile(p, 'r') as zf:
                        if 'AndroidManifest.xml' in zf.namelist():
                            sample_file_name = f'{os.path.basename(p)}.apk'
                            shutil.move(p, os.path.join(self.host_dir, sample_file_name))
                            break

            if not sample_file_name:
                raise DeviceError('APK file not found after extraction.')

            return sample_md5, event_id, event_uuid, package_name
        except Exception as e:
            traceback.print_exc()
            raise DeviceError(f'Error in download_one_sample: {e}')
        finally:
            self.sample_count[0] -= 1

    def clean_host_dir(self):
        def remove_all_files(directory):
            for p in glob.glob(f'{directory}/*'):
                if os.path.isdir(p):
                    shutil.rmtree(p)
                else:
                    os.remove(p)
        if isinstance(self.device, DockerDevice):
            # host_dir is mapped to super_host_dir.``
            remove_all_files(self.device.super_host_dir)
        logger.info('cleaned host directory successfully.')

    def remove_apk(self, sample_md5: str):
        apk_path = f'{self.host_dir}/{sample_md5}.apk'
        os.remove(apk_path)

    def zip_screenshots(self, sample_md5: str):
        """Compress screenshots along with context into a zip archive."""
        files = glob.glob(f'{self.host_dir}/*')
        with zipfile.ZipFile(f'{self.host_dir}/screenshots_{sample_md5}.zip',
                             'w') as zipf:
            for f in files:
                zipf.write(f, arcname=os.path.split(f)[-1])

    def upload_screenshots(self, event_id: int, sample_md5: str):
        """Upload screenshots zip file to misp."""
        attr = pymisp.MISPAttribute()
        attr.type = 'attachment'
        attr.value = f'screenshots_{sample_md5}.zip'
        attr.data = Path(f'{self.host_dir}/screenshots_{sample_md5}.zip')
        a = self.misp.add_attribute(event_id, attr)
        if 'errors' in a.keys():
            return
        self.misp.tag(a['Attribute']['uuid'],
                      'secheart:file-type="screenshots"')
        logger.info('uploaded screenshots successfully.')

    def upload_strace(self, sample_md5: str, event_id: int, event_uuid: str):
        """Upload strace zip file to misp."""
        attr = pymisp.MISPAttribute()
        attr.type = 'attachment'
        attr.value = f'{sample_md5}.strace.zip'
        attr.data = Path(f'{self.host_dir}/{sample_md5}.strace.zip')
        a = self.misp.add_attribute(event_id, attr)
        if 'errors' in a.keys():
            return
        self.misp.tag(a['Attribute']['uuid'], 'secheart:file-type="syscall"')
        self.misp.tag(event_uuid, 'secheart:analysis="syscall-traced"')
        os.remove(f'{self.host_dir}/{sample_md5}.strace.zip')
        os.remove(f'{self.host_dir}/{sample_md5}.strace')
        logger.info('uploaded syscall trace successfully.')

    def upload_pcap(self, sample_md5: str, event_id: int, event_uuid: str):
        """
        Upload pcap file to misp.
        :param event_uuid: uuid of misp event
        :param sample_md5: md5 of the sample.
        :param event_id: misp event id.
        :return:
        """
        pacap_path = f'{self.host_dir}/{sample_md5}.pcap'
        if os.path.getsize(pacap_path) < 100:
            logger.info('pcap file is too small, ignores.')
            os.remove(pacap_path)
            return
        attr = pymisp.MISPAttribute()
        attr.type = 'attachment'
        attr.value = f'{sample_md5}.pcap'
        attr.data = Path(pacap_path)
        a = self.misp.add_attribute(event_id, attr)
        if 'errors' in a.keys():
            return
        self.misp.tag(a['Attribute']['uuid'], 'secheart:file-type="pcap"')
        self.misp.tag(event_uuid,
                      'secheart:analysis="traffic-captured"')
        os.remove(pacap_path)
        logger.info('uploaded pcap successfully.')

    def run(self):
        search_failure_count = 0

        def renew_device():
            try:
                self.device = self.device.renew()
                assert self.device
                self.update_host_dir()
                logger.info(f'renewed device successfully.')

            except Exception as de:
                logger.critical(f'failed to renew device')
                logger.error(traceback.format_exc())
                raise de

        def restart_and_renew_device():
            try:
                self.device.restart()
            except DeviceError:
                renew_device()
            except Exception as de:
                raise de

        while self.sample_count[0] > 0:
            try:
                if not self.device.ensure_availability():
                    renew_device()
                self.clean_host_dir()
                sample_attrs = self.download_one_sample()
                # No eligible sample found anymore
                if not sample_attrs:
                    search_failure_count += 1
                    logger.error('failed to download sample')
                    if search_failure_count > 10:
                        logger.info('no eligible sample found anymore.')
                        self.sample_count[0] = 0
                        break
                    else:
                        continue
                else:
                    search_failure_count = 0

                sample_md5, event_id, event_uuid, package_name = sample_attrs
                self.device.test_app(sample_md5, package_name)
                self.remove_apk(sample_md5)
                try:
                    self.upload_strace(sample_md5, event_id, event_uuid)
                except:
                    logger.info('strace fail.')
                self.upload_pcap(sample_md5, event_id, event_uuid)
                self.zip_screenshots(sample_md5)
                self.upload_screenshots(event_id, sample_md5)
                self.misp.tag(event_uuid, 'secheart:analysis="ui-tested"')
                self.misp.untag(event_uuid, 'secheart:analysis="ui-testing"')
                self.misp.untag(event_uuid, 'secheart:analysis="fail-apk"')
            except DeviceError as e:
                logger.error(str(e).strip("'"))
                if 'failed to built WebDriver' in str(e):
                    continue
                elif 'failed to capture any screenshot' in str(e):
                    continue
                elif 'failed to install' in str(e):
                    continue
                elif 'failed to download a sample' in str(e):
                    continue
                elif 'failed to start the app' in str(e):
                    continue
                elif 'app is not installed' in str(e):
                    continue
                elif 'failed to run adbd as root' in str(e):
                    restart_and_renew_device()
                elif 'failed to start tcpdump' in str(e):
                    restart_and_renew_device()
                elif 'failed to list third party packages' in str(e):
                    restart_and_renew_device()
                elif 'failed to uninstall third party package' in str(e):
                    restart_and_renew_device()
                elif 'failed to make sample data directory' in str(e):
                    restart_and_renew_device()
                elif 'failed to stop tcpdump' in str(e):
                    restart_and_renew_device()
                elif 'failed to remove sample data' in str(e):
                    restart_and_renew_device()
                elif 'failed to restart' in str(e):
                    renew_device()
                else:
                    logger.error(f"can't handle {str(e)}, thread exits.")
                    break
            except ProtocolError as e:
                logger.error(str(e))
                continue
            except Exception as e:
                logger.error(e)
                logger.debug(traceback.format_exc())
                continue
            finally:
                try:
                    self.device.webdriver.quit()
                except Exception as e:
                    logger.debug(e)
                time.sleep(CFG['interval_between_apps'])

 
class AndroidManager:
    """Android sandbox manager."""
    def __init__(self, misp_tags: Union[Sequence[str], str, None], num: int):
        """
        :param misp_tags: MISP tags.
        :param num: number of samples to be tested.
        """
        self.devices: Devices = Devices()
        if isinstance(misp_tags, list):
            self.misp_tags = misp_tags
        else:
            self.misp_tags = [misp_tags]
        # wrap count in list for reference in threads
        self.sample_count: List[int] = [num]
        self.download_lock = Lock()

    def execute(self):
        threads = []
        while not self.devices.empty():
            threads.append(DeviceThread(self.sample_count,
                                        self.devices.get_device(),
                                        self.misp_tags,
                                        self.download_lock))
        for thread in threads:
            thread.start()
            time.sleep(5)
            logger.info(f'{thread.name} is started.')
        for thread in threads:
            thread.join()
        if self.sample_count[0] == 0:
            logger.info('all Apps are tested, done.')
        else:
            logger.critical('Exit unexpectedly.')